package cn.jbolt.ai.resource.embedding.service.impl;

import cn.hutool.core.lang.Snowflake;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.core.embedding.EmbeddingService;
import cn.jbolt.ai.resource.llm.vo.ResourceSelectVO;
import cn.jbolt.util.JSONUtil;
import cn.jbolt.util.Pager;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import cn.jbolt.ai.resource.embedding.EmbeddingModelType;
import cn.jbolt.ai.resource.embedding.dao.EmbeddingResourceMapper;
import cn.jbolt.ai.resource.embedding.entity.EmbeddingResource;
import cn.jbolt.ai.resource.embedding.service.EmbeddingResourceService;
import cn.jbolt.util.Result;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.jboltai.capability.embedding.local.OllamaEmbeddingImpl;
import com.jboltai.resource.embedding.EmbeddingModel;
import com.jboltai.resource.embedding.EmbeddingResourceCenter;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

@Service
public class EmbeddingResourceServiceImpl extends ServiceImpl<EmbeddingResourceMapper, EmbeddingResource> implements EmbeddingResourceService {

    private static final Logger logger = LoggerFactory.getLogger(EmbeddingResourceServiceImpl.class);
    @Autowired
    private EmbeddingService embeddingService;

    @Override
    public Result getResourceById(Long id) {
        EmbeddingResource resource = getById(id);
        if (resource != null) {
            return Result.success(resource);
        } else {
            return Result.error("资源不存在");
        }
    }

    @Override
    public Result listByPage(Integer pageNumber, Integer pageSize, String keywords) {
        QueryWrapper<EmbeddingResource> queryWrapper = new QueryWrapper<>();
        if (StrUtil.isNotEmpty(keywords)) {
            queryWrapper.and(wrapper -> wrapper.like("name", keywords).
                    or().
                    like("type", keywords));
        }
        queryWrapper.orderByDesc("id");
        return Result.success(Pager.from(this.page(new Page<>(pageNumber, pageSize), queryWrapper)));
    }

    @Override
    public Result listAllEnabledType() {

// 主查询，根据子查询的结果过滤原始表中的记录
        QueryWrapper<EmbeddingResource> mainQueryWrapper = new QueryWrapper<>();
        mainQueryWrapper.select("type", "default_model")
                .eq("enable", true).groupBy("type", "default_model");

        List<EmbeddingResource> resourceList = this.list(mainQueryWrapper);

        if (resourceList.size() > 0) {
            List<ResourceSelectVO> selectList = convertToTypeOptions(resourceList);
            return Result.success(selectList);
        } else {
            return Result.success();
        }


    }


    @Override
    public Result listAllEnabledTypeWithTencent() {
        List<ResourceSelectVO> selectList = convertToTypeOptions(this.list(new QueryWrapper<EmbeddingResource>().
                eq("enable", true).orderByDesc("id").
                groupBy("type", "default_model")));

        selectList.add(new ResourceSelectVO(EmbeddingModel.TENCENT_BGE_BASE_ZH.getDefaultModel(), EmbeddingModel.TENCENT_BGE_BASE_ZH.getName()));
        selectList.add(new ResourceSelectVO(EmbeddingModel.TENCENT_M3E_BASE.getDefaultModel(), EmbeddingModel.TENCENT_M3E_BASE.getName()));
        selectList.add(new ResourceSelectVO(EmbeddingModel.TENCENT_E5_LARGE_V2.getDefaultModel(), EmbeddingModel.TENCENT_E5_LARGE_V2.getName()));
        selectList.add(new ResourceSelectVO(EmbeddingModel.TENCENT_TEXT2VEC_LARGE_CHINESE.getDefaultModel(), EmbeddingModel.TENCENT_TEXT2VEC_LARGE_CHINESE.getName()));
        selectList.add(new ResourceSelectVO(EmbeddingModel.TENCENT_MULTILINGUAL_E5_BASE.getDefaultModel(), EmbeddingModel.TENCENT_MULTILINGUAL_E5_BASE.getName()));

        return Result.success(selectList);

    }

    @Override
    @Transactional
    public Result add(EmbeddingResource resource) {
        //ollama模型所有参数都是为必填项
        if (resource.getType().equals("ollama")) {
            if (StrUtil.isBlankIfStr(resource.getDefaultModel())) return Result.error("ollama默认模型不能为空");
            if (StrUtil.isBlankIfStr(resource.getApiUrl())) return Result.error("ollama模型的API_URL不能为空");
            if (resource.getDimension() == null || resource.getDimension() <= 0)
                return Result.error("ollama模型的向量化维度不能为空");
        }
        //其他模型暂时不需要校验
        resource.setEnable(true);
        if (save(resource)) {
            //注册资源
            Result result = registerResource(resource);
            if (result.isError()) {
//                throw new RuntimeException("资源注册失败");
                return Result.error("资源注册失败");

            }
            return Result.success("添加成功");
        } else {
            return Result.error("添加失败");
        }
    }

    @Override
    @Transactional
    public Result update(EmbeddingResource resource) {
        if (updateById(resource)) {
            if (resource.getEnable()) {
                unRegisterResource(resource);
                Result result = registerResource(resource);
                if (result.isError()) {
//                    throw new RuntimeException("资源注册失败");
                    return Result.error("资源注册失败");
                }
            }
            return Result.success();
        } else {
            return Result.error("更新失败");
        }
    }

    @Override
    @Transactional
    public Result del(Long id) {
        Result result = getResourceById(id);
        if (result.isError()) {
            return result;
        }
        EmbeddingResource resource = (EmbeddingResource) result.getData();
        if (resource.getEnable()) {
            unRegisterResource(resource);
        }
        if (removeById(id)) {
            return Result.success();
        } else {
            return Result.error("删除失败");
        }
    }

    @Override
    @Transactional
    public Result toggleEnable(Long id) {
        EmbeddingResource resource = getById(id);
        if (resource == null) {
            return Result.error("资源不存在");
        }
        Boolean enable = !resource.getEnable();
        resource.setEnable(enable);
        boolean updated = updateById(resource);
        if (updated) {
            if (enable) { //开启的时候注册资源
                Result result = registerResource(resource);
                if (result.isError()) {
//                    throw new RuntimeException("资源注册失败");
                    return Result.error("资源注册失败");
                }
                return Result.success();
            } else {//关闭的时候移除资源
                unRegisterResource(resource);
                return Result.success();
            }

        } else {
            return Result.error("更新启用状态失败");
        }
    }

    @Override
    public Result getModelDimension(EmbeddingResource embeddingResource) {
        if (StrUtil.isBlankIfStr(embeddingResource.getDefaultModel())) {
            return Result.error("ollama默认模型不能为空");
        }
        if (StrUtil.isBlankIfStr(embeddingResource.getApiUrl())) {
            return Result.error("ollama模型的API_URL不能为空");
        }

        //创建一个临时id
        Snowflake snowflake = IdUtil.getSnowflake();
        long id = snowflake.nextId();
        embeddingResource.setId(id);
        embeddingResource.setDimension(0);
        Result finalResult = null;
        try {
            finalResult = registerResource(embeddingResource);
            if (finalResult.isSuccess()) {
                finalResult = embeddingService.getModelDimension(
                        EmbeddingModel.getCustomModel("ollama-" + embeddingResource.getDefaultModel()), embeddingResource);
            }
        } catch (Exception e) {
            logger.error("Embedding资源--{}-{}---获取维度失败,异常原因:{}", embeddingResource.getType(), embeddingResource.getDefaultModel(), e.getMessage());
            return Result.error(e.getMessage());
        } finally {
            if (finalResult.isSuccess()) {
                unRegisterResource(embeddingResource);
            }

        }
        return finalResult;
    }

    /**
     * 移除资源
     *
     * @param resource 资源
     */
    public void unRegisterResource(EmbeddingResource resource) {
        logger.info("Embedding资源--{}-{}---移除资源", resource.getType(), resource.getDefaultModel());

        //从资源池
        try {
            com.jboltai.resource.embedding.EmbeddingResource jboltEmbeddingResource = EmbeddingResourceCenter.getResourceById(resource.getId().toString());
            EmbeddingResourceCenter.unregister(jboltEmbeddingResource);
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("Embedding资源--{}-{}---移除资源失败,异常原因:{}", resource.getType(), resource.getDefaultModel(), e.getMessage());
            throw e;
        }
    }

    /**
     * 注册所有资源,是在项目启动时候调用
     */
    public void registerAllResource() {
        //只启动开启状态的资源
        this.list(new QueryWrapper<EmbeddingResource>().eq("enable", true)).
                forEach(this::registerResource);
    }

    /**
     * 注册资源
     *
     * @param embeddingResource 资源
     */
    public Result registerResource(@NotNull EmbeddingResource embeddingResource) {

        String type = embeddingResource.getType();
        String defaultModel = embeddingResource.getDefaultModel();
        String apiUrl = embeddingResource.getApiUrl();
        String apiKey = embeddingResource.getApiKey();
        String apiSecret = embeddingResource.getApiSecret();
        Integer dimension = embeddingResource.getDimension();
        logger.info("Embedding资源-id:{}-{}--注册资源", embeddingResource.getId(), type + (type.equals("ollama") == true ? ("-" + defaultModel) : ""));
        //创建模型
        Result modelResult = getOrCreateModel(type, defaultModel, dimension);
        if (modelResult.isError()) return modelResult;

        //创建资源
        com.jboltai.resource.embedding.EmbeddingResource jboltEmbeddingResource =
                new com.jboltai.resource.embedding.EmbeddingResource((EmbeddingModel) modelResult.getData(),
                        embeddingResource.getId().toString(), StrUtil.isNotEmpty(apiKey) ? apiKey : null);
        if (StrUtil.isNotEmpty(defaultModel)) jboltEmbeddingResource.setDefaultModelName(defaultModel);
        if (StrUtil.isNotEmpty(apiUrl)) jboltEmbeddingResource.setApiUrl(apiUrl);
        if (StrUtil.isNotEmpty(apiSecret)) jboltEmbeddingResource.setAppSecret(apiSecret);


        setRequestLimits(embeddingResource.getLimits(), jboltEmbeddingResource);
        //注册资源
        try {
            EmbeddingResourceCenter.register(jboltEmbeddingResource);
        } catch (Exception e) {
            logger.error("Embedding资源--{}--注册资源失败,异常原因:{}", type + (type.equals("ollama") == true ? ("-" + defaultModel) : ""), e.getMessage());
            return Result.error(e.getMessage());
        }
        return Result.success();
    }

    /**
     * 获取模型
     * ollama跟其他模型不一样,需要自定义模型,
     * 其他模型可以直接获取
     *
     * @param type
     * @param defaultModel
     * @param dimension
     * @return
     */
    public Result getOrCreateModel(String type, String defaultModel, Integer dimension) {

        if (type.equals("ollama")) {
            String modelName = "ollama-" + defaultModel;//ollama模型名称的名称是type+defaultModel自定义的
            EmbeddingModel model = EmbeddingModel.getCustomModel(modelName);
            if (model == null) {
//                model = new EmbeddingModel(type, modelName, dimension, OllamaEmbeddingImpl.INSTANCE, 9999, 1);
                model = new EmbeddingModel(modelName, defaultModel, dimension, OllamaEmbeddingImpl.INSTANCE, 9999, 1);
                //在这里加上一个trycath块,防止自定义模型已经存在了,导致报错
                try {
                    EmbeddingModel.registerCustomModel(model);
                } catch (Exception e) {
                    String message = e.getMessage();
                    logger.error("ollama自定义模型异常,异常原因:{}", message);
                    return Result.error(message);
                }

            }
            return Result.success(model);
        } else { //除了ollama模型,其他模型都可以通过getModel方法获取
            EmbeddingModel model = EmbeddingModel.getModel(type);
            return Result.success(model);
        }

    }

    /**
     * 统一设置频率限制
     *
     * @param limitsStr
     * @param jboltEmbeddingResource
     */

    private void setRequestLimits(String limitsStr, com.jboltai.resource.embedding.EmbeddingResource jboltEmbeddingResource) {
        if (StrUtil.isNotEmpty(limitsStr) && JSONUtil.isTypeJSONArray(limitsStr)) {
            ArrayNode jsonArray = JSONUtil.toJsonArray(limitsStr);
            for (int i = 0; i < jsonArray.size(); i++) {
                JsonNode jsonObject = jsonArray.get(i);
                String unit = jsonObject.get("unit").asText();
                int times = jsonObject.get("times").asInt();
                if (times > 0) {
                    switch (unit.toLowerCase()) {
                        case "minute":
                            jboltEmbeddingResource.setRequestLimit(TimeUnit.MINUTES, times);
                            break;
                        case "hour":
                            jboltEmbeddingResource.setRequestLimit(TimeUnit.HOURS, times);
                            break;
                        case "day":
                            jboltEmbeddingResource.setRequestLimit(TimeUnit.DAYS, times);
                            break;

                    }
                }
            }
        }
    }

    /**
     * 根据类型转换成下拉选列表
     *
     * @param list
     * @return
     */
    public List<ResourceSelectVO> convertToTypeOptions(List<EmbeddingResource> list) {
        List<ResourceSelectVO> selectList = new ArrayList<>();
        for (EmbeddingResource resource : list) {
            ResourceSelectVO selectVO = new ResourceSelectVO();
            selectVO.setSn(resource.getType() + ":" + (StrUtil.isNotEmpty(resource.getDefaultModel()) ? resource.getDefaultModel() : ""));
            if (resource.getType().equals(EmbeddingModelType.OLLAMA.getSn())) { //ollama模型特殊,名称是type-defaultModel组成的
                selectVO.setName(EmbeddingModelType.getBySn(resource.getType()).getName() + "-" + resource.getDefaultModel());
            } else {
                selectVO.setName(EmbeddingModelType.getBySn(resource.getType()).getName());
            }
            selectList.add(selectVO);
        }
        return selectList;
    }

    /**
     * 根据类型和默认模型获取资源
     *
     * @param type
     * @param defaultModel
     * @return
     */
    public EmbeddingResource getResourceByTypeAndDefaultModel(String type, String defaultModel) {
        return this.getOne(new QueryWrapper<EmbeddingResource>().eq("type", type).eq("default_model", defaultModel).last("limit 1"));
    }
}
